import os
import time
from typing import Any, Union, Callable

import torch
import torchvision
from torch.utils import data
from torch.utils.data import DataLoader
from torchvision import transforms
from d2l import torch as d2l
from matplotlib import pyplot as plt


def load_data_fashion_mnist(batch_size: int, cpu_workers: int, resize=None, downloaded=True) -> tuple[
    DataLoader[Any], DataLoader[Any]]:
    trans = [transforms.ToTensor()]  ## 由于可能会有多个预处理步骤，Compose函数在接受预处理流程时需要接收一个列表传入
    if resize:
        trans.insert(0, transforms.Resize(resize))
    trans = transforms.Compose(trans)
    mnist_train = torchvision.datasets.FashionMNIST(
        root='./data',
        transform=trans,
        train=True,
        download=(not downloaded)
    )
    mnist_test = torchvision.datasets.FashionMNIST(
        root='./data',
        transform=trans,
        train=False,
        download=(not downloaded)
    )
    return (
        data.DataLoader(mnist_train, batch_size, shuffle=True, num_workers=cpu_workers),
        data.DataLoader(mnist_test, batch_size, shuffle=False, num_workers=cpu_workers)
    )


class Accumulator:
    def __init__(self, n):
        self.data = [0.0] * n  ##创建n个计数器

    def add(self, *args):
        self.data = [a + float(b) for a, b in zip(self.data, args)]

    def __getitem__(self, idx):
        return self.data[idx]


def accuracy(y_predict: torch.Tensor, y: torch.Tensor) -> torch.Tensor:
    if len(y_predict.shape) > 1 and y_predict.shape[1] > 1:
        y_predict = y_predict.argmax(axis=1)  ## 这里是获取最大元素的索引，不是值！！
    cmp = y_predict.type(y.dtype) == y
    cmp.type(y.dtype)  ## 将cmp的数据格式转为y的数据格式，未转换前是bool，转换后为integer
    return float(cmp.sum())


def evaluate_accuracy(net: Union[Callable, torch.nn.Module], data_iter: torch.utils.data.DataLoader):
    if isinstance(net, torch.nn.Module):  ##判断net是否为torch.nn.Module
        net.eval()
    metric = Accumulator(2)
    for X, y in data_iter:
        metric.add(
            accuracy(
                net(X), y
            ),
            y.numel()
        )
    return metric[0] / metric[1]
